package org.example;

import com.google.common.collect.Lists;
import org.junit.jupiter.api.Test;

import java.util.*;
import java.util.stream.Collectors;

public class Backtrack {

    @Test
    public void run() {
        //        List<Integer> choices = List.of(1, 2, 2);
        List<Integer> choices = List.of(1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0);

        List<List<Integer>> res = new ArrayList<>();
        boolean[] selected = new boolean[choices.size()];
        backtrack(new ArrayList<>(), choices, selected, res);
        System.out.println(res.size());
        for (List<Integer> re : res) {
            System.out.println(re.stream().map(String::valueOf).collect(Collectors.joining()));
        }
    }

    /**
     * 有序回溯
     * @param state
     * @param choices
     * @param selected
     * @param res
     */
    private void backtrack(List<Integer> state,
                           List<Integer> choices,
                           boolean[] selected,
                           List<List<Integer>> res) {
        if (state.size() == choices.size()) {
            res.add(new ArrayList<>(state));
            return;
        }

        // 重复剪枝 基于同层
        Set<Integer> dup = new HashSet<>();
        // 存在分支 使用循环
        for (int i = 0; i < choices.size(); i++) {
            Integer choice = choices.get(i);
            // 当前节点标志 需要全链路传播
            if (selected[i]) {
                continue;
            }

            if (dup.contains(choice)) {
                continue;
            }
            dup.add(choice);

            state.add(choice);
            selected[i] = true;
            backtrack(state, choices, selected, res);
            state.remove(state.size() - 1);
            selected[i] = false;
        }
    }

    @Test
    void runV4() {
        List<Integer> choices = List.of(1, 2, 3);
        int target = 5;
        int[] dp = new int[target + 1];
        dp[1] = 1;
        dp[2] = 2;
        dp[3] = 4;

        for (int i = 4; i <= target; i++) {
            dp[i] = dp[i - 1] + dp[i - 2] + dp[i - 3];
        }

        System.out.println(22);

    }

    @Test
    void runV3() {
        List<Integer> choices = List.of(1,2, 3);
        int target = 5;
        List<List<Integer>> res = new ArrayList<>();
        backtrackV3(new ArrayList<>(), choices, 0, target, res);

        for (List<Integer> re : res) {
            System.out.println(re.stream().map(String::valueOf).collect(Collectors.joining()));
        }

        System.out.println(res.size());
    }

    /**
     * 无序回溯
     * @param state
     * @param choices
     * @param start
     * @param target
     * @param res
     */
    void backtrackV3(List<Integer> state,
                     List<Integer> choices,
                     int start,
                     int target,
                     List<List<Integer>> res) {

//        start = 0;
        // 计算当前值
        if (state.stream().mapToInt(Integer::intValue).sum() == target) {
            res.add(new ArrayList<>(state));
            return;
        }

        for (int i = start; i < choices.size(); i++) {
            if (state.stream().mapToInt(Integer::intValue).sum() > target) {
                return;
            }

            Integer choice = choices.get(i);
            state.add(choice);
            backtrackV3(state, choices, i, target, res);
            state.remove(choice);
        }
    }

    /* 回溯算法：n 皇后 */
    void backtrack(int row,
                   int n,
                   List<List<String>> state,
                   List<List<List<String>>> res,
                   boolean[] cols,
                   boolean[] diags1,
                   boolean[] diags2) {

        if (row == n) {
            List<List<String>> copyState = new ArrayList<>();
            for (List<String> sRow : state) {
                copyState.add(new ArrayList<>(sRow));
            }
            res.add(copyState);
            return;
        }

        for (int col = 0; col < n; col++) {
            int idx1 = row - col + n - 1;
            int idx2 = row + col;

            if (cols[col] || diags1[idx1] || diags2[idx2]) {
                continue;
            }

            state.get(row).set(col, "Q");
            cols[col] = diags2[idx2] = diags1[idx1] = true;
            backtrack(row + 1, n, state, res, cols, diags1, diags2);
            state.get(row).set(col, "#");
            cols[col] = diags2[idx2] = diags1[idx1] = false;
        }
    }

    /* 求解 n 皇后 */
    List<List<List<String>>> nQueens(int n) {
        // 初始化 n*n 大小的棋盘，其中 'Q' 代表皇后，'#' 代表空位
        List<List<String>> state = new ArrayList<>();
        for (int i = 0; i < n; i++) {
            List<String> row = new ArrayList<>();
            for (int j = 0; j < n; j++) {
                row.add("#");
            }
            state.add(row);
        }
        boolean[] cols = new boolean[n]; // 记录列是否有皇后
        boolean[] diags1 = new boolean[2 * n - 1]; // 记录主对角线上是否有皇后
        boolean[] diags2 = new boolean[2 * n - 1]; // 记录次对角线上是否有皇后
        List<List<List<String>>> res = new ArrayList<>();

        backtrack(0, n, state, res, cols, diags1, diags2);

        return res;
    }

    @Test
    void test() {

        List<List<List<String>>> lists = nQueens(4);
        for (List<List<String>> list : lists) {
            System.out.println(list);
        }
    }
}
